The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

I think we need to have a low_cpu_mem_usage tag in the load_adapter() method too: https://github.com/huggingface/diffusers/blob/28f9d84549c0b1d24ef00d69a4c723f3a11cffb6/src/diffusers/loaders/lora_pipeline.py#L371C30-L371C42

If so, then this PR would be contingent on that. We could, however, use a combination of inject_adapter_in_model() and set_peft_model_state_dict() to mimic the same thing, I assume. I wouldn't personally prefer that because load_adapter() has been there for a while in diffusers.

Could you please clarify that bit? On the PEFT side, we have low_cpu_mem_usage on load_adapter but that's not the method being used here (just has the same name), right? Is this method coming from transformers (i.e. here)?

that's not the method being used here (just has the same name), right? Is this method coming from transformers (i.e. here)?

Yes, that is correct.

@BenjaminBossan when I used:

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to("cuda")
pipe.load_lora_weights(
    "TheLastBen/The_Hound", 
    weight_name="sandor_clegane_single_layer.safetensors", 
    low_cpu_mem_usage=True
)

prompt = "sandor clegane drinking in a pub"
image = pipe(
    prompt=prompt,
    num_inference_steps=30,
    width=1024,
    generator=torch.manual_seed(42),
    height=1024,
).images[0]
image.save("sandor.png")

It leads to:

  File "/home/sayak/diffusers/src/diffusers/models/transformers/transformer_flux.py", line 98, in forward
    hidden_states = gate * self.proj_out(hidden_states)
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/peft/tuners/lora/layer.py", line 585, in forward
    result = result + lora_B(lora_A(dropout(x))) * scaling
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/sayak/.pyenv/versions/3.10.12/envs/diffusers/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 125, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper_CUDA_mm)

I investigated this a bit and confirmed that the LoRA params are kept on CPU which causes this failure. In case of low_cpu_mem_usage=False the LoRA parameters are on the expected device ("cuda" in the above example).

I further investigated why the tests added in this PR don't fail. That is because the state dict we're supplying to set_peft_model_state_dict() (here) -- the tensors of that state dict are already on the desired device. When I forcibly changed their device to a CPU and ran the tests on a GPU, the tests failed, and they complained about the same thing.

Possible to look into this?

Hmm, I could not reproduce the issue :-/ I had to change the code slightly due to memory constraints:

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, device_map="balanced", max_memory={0: "24GB", 1: "20GB"},
)
# after loading the LoRA adapter {p.device for p in pipe.transformer.parameters()} returns:
# {device(type='cuda', index=0)}

Could this be the reason why it works for me?

I also tried this with a normal PEFT model that I moved to CUDA and then loaded with low_cpu_mem_usage and it worked.

Yeah you need to be on the exact same setup to replicate this. We cannot assume people will do load_lora_weights() only in a specific manner.

You can perhaps use an SD LoRA:

Quick update, I couldn't reproduce with that model:

import torch
from diffusers import FluxPipeline, StableDiffusionPipeline

generator = torch.Generator().manual_seed(0)
pipe = StableDiffusionPipeline.from_pretrained("hf-internal-testing/Counterfeit-V2.5", safety_checker=None).to("cuda")
lora_model_id = "hf-internal-testing/civitai-light-shadow-lora"
lora_filename = "light_and_shadow.safetensors"
pipe.load_lora_weights(lora_model_id, weight_name=lora_filename, low_cpu_mem_usage=True)
images = pipe(
    "masterpiece, best quality, mountain", output_type="np", generator=generator, num_inference_steps=2
).images

@BenjaminBossan here's a minimal reproduction:

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("sayakpaul/tiny-flux-pipeline-with-lora", torch_dtype=torch.bfloat16).to("cuda")
pipe.load_lora_weights(
    "sayakpaul/tiny-flux-pipeline-with-lora", weight_name="pytorch_lora_weights.bin", low_cpu_mem_usage=True
)

prompt = "sandor clegane drinking in a pub"
image = pipe(prompt=prompt, num_inference_steps=30).images[0]
image.save("sandor.png")

I am on peft:main.

Okay, got it now, thanks for the memory-friendly reproducer.

Indeed, if the LoRA weights on the model are on meta device, the device will be taken from the state_dict, not the base layer. I worked on a fix: huggingface/peft#2113

For the time being, you could add this snippet and it should fix the issue:

    if low_cpu_mem_usage:
        for module in model.modules():
            if hasattr(module, "_move_adapter_to_device_of_base_layer"):
                module._move_adapter_to_device_of_base_layer(adapter_name)

@BenjaminBossan could you give this a review?

@BenjaminBossan thanks!

Docstrings still contain TODO

Yeah will be resolved after Yiyi's approval.

Since PEFT version v0.13.1 is now released, the min PEFT version should be updated accordingly.

That's done.

In case you plan on no longer supporting older PEFT and transformer versions in the future: I would add a TODO comment to all those version checks that they can be removed once support for those older versions is dropped. If you plan to support them indefinitely, ignore this comment.

For now, we can ignore it.

An entry to the diffusers PEFT docs would be nice to have, especially since the name of the argument is not really intuitive.

Done.

@yiyixuxu could you review this PR?

After the approval, will add docs and request for a review from Steven.

Thanks, @yiyixuxu!

I have taken care of the TODOs in the docs too. @stevhliu could you review the related changes?